The metrics should not care if the input is low memory or not, they should only care about volumes1, volumes2, and some arguments if necessary. The whole low memory stuff should come at the pipeline level. Here is what I would do

1. Implement functions that normalize volumes the way you want it, they should take a stack of volumes or a single volumes. Don't do this things inside an if-statement.

2. Load a chunk of the gt volumes, run all the preprocessing you need, and then get the submatrix of the full distance matrix for that chunk of volumes. This should use the distance matrices you already have implemented. Populate the full distance matrix on-the-fly. This way you can do a vmap to compute a submatrix, and then a simple stack operation to get the full matrix.

I would also create a low_memory_pipeline and a regular memory pipeline. All those if statements make the code very messy, and I think it'd be easy to miss something being wrong.

Yes, good point! I forgot that the getitem takes care of indexing into sub-batches

I did the masking and normalization in the subbatch. This prevents a duplication of reading it into memory (once for normalization / masking, another right before doing the m2m distance call)

There is still some if low memory mode control flow. To reduce code duplication.

Also, now the gt maps should be a .npy file, not a .pt file.

I hope that torch.from_numpy(np.load(fname.npy)) is not a lot slower than torch.load(fname.pt) for the big 160 GB file.